from manimlib import *
import numpy as np
from numpy.polynomial.chebyshev import Chebyshev


# 定义待逼近的复杂函数
def func(x):
    return np.sin(2 * np.pi * x) + 0.5 * np.cos(3 * np.pi * x) + x**2


# 多项式做基函数
class PolyGalerkin(Scene):
    def construct(self):
        # 基函数
        def basis_func(i, x):
            return x**i
        
        # 逼近函数
        def appro_func(x, params=[]):
            res = 0
            for i, c in enumerate(params):
                res += basis_func(i, x) * c
            return res
        
        axes = Axes(
            # x-axis ranges from -1 to 10, with a default step size of 1
            x_range=(0, 1, 0.1),
            # y-axis ranges from -2 to 2 with a step size of 0.5
            y_range=(-1, 1, 1),
            # The axes will be stretched so as to match the specified
            # height and width
            height=6,
            width=10,
            # Axes is made of two NumberLine mobjects.  You can specify
            # their configuration with axis_config
            axis_config=dict(
                stroke_color=GREY_A,
                stroke_width=2,
                numbers_to_exclude=[0],
            ),
            # Alternatively, you can specify configuration for just one
            # of them, like this.
            y_axis_config=dict(
                # big_tick_numbers=[-2, 2],
            ),
        )
        axes.add_coordinate_labels(
            font_size=20,
            num_decimal_places=1,
        )

        self.play(Write(axes, lag_ratio=0.01, run_time=1))

        # Axes.get_graph will return the graph of a function
        func_graph = axes.get_graph(func, color=RED)
        func_label = axes.get_graph_label(func_graph, "function: f(x)=\\sin(2\\pi x)+0.5\\cos(3\\pi x)+x^2")
        func_label.next_to(func_graph, DOWN, buff=1.5)

        self.play(
            ShowCreation(func_graph),
            FadeIn(func_label, DOWN),
            run_time=1,
        )

        param_ans = [0.7034375445786613, 1.7835206081854447, -8.571187092028353, 6.380791394685604, 0]

        appro_graph = axes.get_graph(lambda x: appro_func(x), color=BLUE)
        appro_label = Tex("g(x)=0")
        appro_label.next_to(func_label, DOWN, buff=0.5)
        self.play(ShowCreation(appro_graph), ShowCreation(appro_label), run_time=1)

        tex_str = "g(x)="
        # 创建动画，改变参数i
        for i, c in enumerate(param_ans):
            if c>0:
                tex_str +="+"
            tex_str += f"{"{:.2f}".format(c)}x^{i}"
            new_appro_label = Tex(tex_str)
            new_appro_label.next_to(func_label, DOWN, buff=0.5)
            if c != 0:
                self.play(Transform(appro_label, new_appro_label), run_time=1)
            for t in np.linspace(0, c, 20):
                new_graph = axes.get_graph(lambda x: appro_func(x, param_ans[: i + 1]), color=BLUE)
                self.play(Transform(appro_graph, new_graph),  run_time=0.1)
        return


# 多项式做基函数
class SineGalerkin(Scene):
    def construct(self):
        # 基函数
        def basis_func(i, x):
            if i == 0:
                return np.sin(np.pi * x)
            else:
                return np.sin((i + 1) * np.pi * x)
        
        # 逼近函数
        def appro_func(x, params=[]):
            res = 0
            for i, c in enumerate(params):
                res += basis_func(i, x) * c
            return res
        
        axes = Axes(
            # x-axis ranges from -1 to 10, with a default step size of 1
            x_range=(0, 1, 0.1),
            # y-axis ranges from -2 to 2 with a step size of 0.5
            y_range=(-1, 1, 1),
            # The axes will be stretched so as to match the specified
            # height and width
            height=6,
            width=10,
            # Axes is made of two NumberLine mobjects.  You can specify
            # their configuration with axis_config
            axis_config=dict(
                stroke_color=GREY_A,
                stroke_width=2,
                numbers_to_exclude=[0],
            ),
            # Alternatively, you can specify configuration for just one
            # of them, like this.
            y_axis_config=dict(
                # big_tick_numbers=[-2, 2],
            ),
        )
        axes.add_coordinate_labels(
            font_size=20,
            num_decimal_places=1,
        )

        self.play(Write(axes, lag_ratio=0.01, run_time=1))

        # Axes.get_graph will return the graph of a function
        func_graph = axes.get_graph(func, color=RED)
        func_label = axes.get_graph_label(func_graph, "function: f(x)=\\sin(2\\pi x)+0.5\\cos(3\\pi x)+x^2")
        func_label.next_to(func_graph, DOWN, buff=1.5)

        self.play(
            ShowCreation(func_graph),
            FadeIn(func_label, DOWN),
            run_time=1,
        )

        param_ans = [0.3786075, 0.4270422, 0.20265058, 0.20462778, 0]

        appro_graph = axes.get_graph(lambda x: appro_func(x), color=BLUE)
        appro_label = Tex("g(x)=0")
        appro_label.next_to(func_label, DOWN, buff=0.5)
        self.play(ShowCreation(appro_graph), ShowCreation(appro_label), run_time=1)

        tex_str = "g(x)="
        basis_str = ["\\sin(\\pi x)",
                     "\\sin(2\\pi x)",
                     "\\sin(3\\pi x)",
                     "\\sin(4\\pi x)",
                     "\\sin(5\\pi x)",
                     ]
        # 创建动画，改变参数i
        for i, c in enumerate(param_ans):
            if c>0:
                tex_str +="+"
            tex_str += f"{"{:.2f}".format(c)}{basis_str[i]}"
            new_appro_label = Tex(tex_str)
            new_appro_label.next_to(func_label, DOWN, buff=0.5)
            if c != 0:
                self.play(Transform(appro_label, new_appro_label), run_time=1)
            for t in np.linspace(0, c, 20):
                new_graph = axes.get_graph(lambda x: appro_func(x, param_ans[: i + 1]), color=BLUE)
                self.play(Transform(appro_graph, new_graph),  run_time=0.1)
        return


# 多项式做基函数
class ChebyshevGalerkin(Scene):
    def construct(self):
        # 基函数
        def basis_func(i, x):
            return Chebyshev.basis(i)(x)
        
        # 逼近函数
        def appro_func(x, params=[]):
            res = 0
            for i, c in enumerate(params):
                res += basis_func(i, x) * c
            return res
        
        axes = Axes(
            # x-axis ranges from -1 to 10, with a default step size of 1
            x_range=(0, 1, 0.1),
            # y-axis ranges from -2 to 2 with a step size of 0.5
            y_range=(-1, 1, 1),
            # The axes will be stretched so as to match the specified
            # height and width
            height=6,
            width=10,
            # Axes is made of two NumberLine mobjects.  You can specify
            # their configuration with axis_config
            axis_config=dict(
                stroke_color=GREY_A,
                stroke_width=2,
                numbers_to_exclude=[0],
            ),
            # Alternatively, you can specify configuration for just one
            # of them, like this.
            y_axis_config=dict(
                # big_tick_numbers=[-2, 2],
            ),
        )
        axes.add_coordinate_labels(
            font_size=20,
            num_decimal_places=1,
        )

        self.play(Write(axes, lag_ratio=0.01, run_time=1))

        # Axes.get_graph will return the graph of a function
        func_graph = axes.get_graph(func, color=RED)
        func_label = axes.get_graph_label(func_graph, "function: f(x)=\\sin(2\\pi x)+0.5\\cos(3\\pi x)+x^2")
        func_label.next_to(func_graph, DOWN, buff=1.5)

        self.play(
            ShowCreation(func_graph),
            FadeIn(func_label, DOWN),
            run_time=1,
        )

        param_ans = [-3.58215600e+00,6.56911415e+00, -4.28559355e+00,1.59519785e+00, 0]

        appro_graph = axes.get_graph(lambda x: appro_func(x), color=BLUE)
        self.play(ShowCreation(appro_graph), run_time=1)

        # 创建动画，改变参数i
        for i, c in enumerate(param_ans):
            if c != 0:
                self.wait(1)
            for t in np.linspace(0, c, 20):
                new_graph = axes.get_graph(lambda x: appro_func(x, param_ans[: i + 1]), color=BLUE)
                self.play(Transform(appro_graph, new_graph),  run_time=0.1)
        return
